#include "owl_math.h"
#include <math.h>
#include "common.h"
#include "logging.h"

void OwlFill(int N, Float alpha, Float *x, int incx) {
	int i = 0;
	for (i = 0; i < N; ++i) {
		x[i * incx] = alpha;
	}
}

void OwlCopy(int N, Float *x, int incx, Float *y, int incy) {
	int i = 0;
	for (i = 0; i < N; ++i) {
		y[i * incy] = x[i * incx];
	}
}

void OwlMean(Float *x, int batch, int filters, int spatial, Float *mean) {
	Float scale = 1.0 / (batch * spatial);
	int i, j, k;
	for (i = 0; i < filters; ++i) {
		mean[i] = 0.0;
		for (j = 0; j < batch; ++j) {
			for (k = 0; k < spatial; ++k) {
				int index = j * filters * spatial + i * spatial + k;
				mean[i] += x[index];
			}
		}
		mean[i] *= scale;
	}
}

void OwlVariance(Float *x, Float *mean, int batch, int filters, int spatial, Float *variance) {
	Float scale = 1.0 / (batch * spatial - 1);
	int i, j, k;
	for (i = 0; i < filters; ++i) {
		variance[i] = 0.0;
		for (j = 0; j < batch; ++j) {
			for (k = 0; k < spatial; ++k) {
				int index = j * filters * spatial + i * spatial + k;
				variance[i] += pow((x[index] - mean[i]), 2);
			}
		}
		variance[i] *= scale;
	}
}

void OwlScal(int N, Float alpha, Float *x, int incx) {
	int i = 0;
	for (i = 0; i < N; ++i) {
		x[i * incx] *= alpha;
	}
}

void OwlNormalize(Float *x, Float *mean, Float *variance, int batch, int filters, int spatial) {
	int b, f, i;
	for (b = 0; b < batch; ++b) {
		for (f = 0; f < filters; ++f) {
			for (i = 0; i < spatial; ++i) {
				int index = b * filters * spatial + f * spatial + i;
				x[index] = (x[index] - mean[f])/(sqrt(variance[f]) + 0.000001f);
			}
		}
	}
}

void OwlAxpy(int N, Float alpha, Float *x, int incx, Float *y, int incy) {
	int i = 0;
	for (i = 0; i < N; ++i) {
		y[i * incy] += alpha * x[i * incx];
	}
}


void OwlGemmNN(int M, int N, int K, Float alpha,
	Float *A, int lda, Float *B, int ldb, Float *C, int ldc) {
	int i, j, k;
	for (i = 0; i < M; ++i) {
		for (k = 0; k < K; ++k) {
			register Float a_part = alpha * A[i * lda + k];
			for (j = 0; j < N; ++j) {
				C[i *ldc + j] += a_part * B[k * ldb + j];
			}
		}
	}
}

void OwlGemmNT(int M, int N, int K, Float alpha,
	Float *A, int lda, Float *B, int ldb, Float *C, int ldc) {
	int i, j, k;
	for (i = 0; i < M; ++i) {
		for (j = 0; j < N; ++j) {
			register Float sum = 0;
			for (k = 0; k < K; ++k) {
				sum += alpha * A[i * lda + k] * B[j * ldb + k];
			}
			C[i * ldc + j] += sum;
		}
	}
}

void OwlGemmTN(int M, int N, int K, Float alpha,
	Float *A, int lda, Float *B, int ldb, Float *C, int ldc) {
	int i, j, k;
	for (i = 0; i < M; ++i) {
		for (k = 0; k < K; ++k) {
			register Float a_part = alpha * A[k * lda + i];
			for (j =0; j < N; ++j) {
				//printf("A[%d][%d]=%lf B[%d][%d]=%lf \n", k, i, A[k * lda + i], k, j,  B[k * ldb + j]);
				C[i * ldc + j] += a_part * B[k * ldb + j];
			}
		}
	}
}


void OwlGemmTT(int M, int N, int K, Float alpha,
	Float *A, int lda, Float *B, int ldb, Float *C, int ldc) {
	int i, j, k;
	for (i = 0; i < M; ++i) {
		for (j = 0; j < N; ++j) {
			register Float sum =0;
			for (k = 0; k < K; ++k) {
				//printf("A[%d][%d]=%lf B[%d][%d]=%lf \n", k, i, A[k * lda + i], j, k,  B[j * ldb + k]);
				sum += alpha * A[k * lda + i] * B[j * ldb + k];
			}
			C[i * ldc + j] += sum;
		}
	}
}

void OwlGemm(int TA, int TB, int M, int N, int K, Float alpha,
	Float *A, int lda, Float *B, int ldb, Float beta, Float *C, int ldc) {
	int i, j;
	for (i = 0; i < M; ++i) {
		for (j = 0; j < N; ++j) {
			C[i * ldc + j] *= beta;
		}
	}

	if (!TA && !TB) {
		//printf("1\n");
		OwlGemmNN(M, N, K, alpha, A, lda, B, ldb, C, ldc);
	} else if (TA && !TB) {
		//printf("2\n");
		OwlGemmTN(M, N, K, alpha, A, lda, B, ldb, C, ldc);
	} else if (!TA && TB) {
		//printf("3\n");
		OwlGemmNT(M, N, K, alpha, A, lda, B, ldb, C, ldc);
	} else {
		//printf("4\n");
		OwlGemmTT(M, N, K, alpha, A, lda, B, ldb, C, ldc);
	}
}

void OwlL2(int n, Float *pred, Float *truth, Float *delta, Float *error) {
	OWLLOG("OwlL2", pred, n);
	OWLLOG("OwlL2, truth ", truth, n);
	int i = 0;
	for (i = 0; i < n; ++i) {
		Float diff = truth[i] - pred[i];
		error[i] = diff * diff;
		delta[i]  = diff;
	}
	OWLLOG("OwlL2, delta ", delta, n);
	OWLLOG("OwlL2, error ", error, n);
}

Float OwlSum(int n, Float *a) {
	int i = 0;
	Float sum = 0.0;
	for (i = 0; i < n; ++i) {
		sum += a[i];
	}
	return sum;
}
